# -*- coding: utf-8 -*-
import jieba
import jieba.analyse
from wordcloud import WordCloud
import matplotlib.pyplot as plt
from collections import Counter
import numpy as np

# 创建情感词云图的函数
def generate_sentiment_wordcloud(sentiment_dict, sentiment_type, title):
    """生成情感词云图"""
    # 准备词频数据
    words_freq = {}
    for word, score in sentiment_dict:
        # 取绝对值作为词频
        words_freq[word] = abs(score)
    
    # 检查是否有词汇
    if not words_freq:
        print(f"没有足够的{sentiment_type}词汇生成词云图")
        return
    
    # 创建词云对象
    wc = WordCloud(
        font_path="simhei.ttf",  # 使用黑体字体，支持中文
        width=800,
        height=600,
        background_color='white',
        max_words=200,
        colormap='viridis' if sentiment_type == 'positive' else 'plasma'
    )
    
    # 生成词云
    wc.generate_from_frequencies(words_freq)
    
    # 显示词云图
    plt.figure(figsize=(10, 8))
    plt.imshow(wc, interpolation='bilinear')
    plt.axis('off')
    plt.title(title, fontsize=20)
    plt.tight_layout()
    
    # 保存词云图
    filename = f"{sentiment_type}_wordcloud.png"
    plt.savefig(filename, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"{sentiment_type}情感词云图已保存为 {filename}")

def load_sentiment_words_from_result(result_file):
    """从情感分析结果文件中加载情感词汇"""
    positive_words = []
    negative_words = []
    
    with open(result_file, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    
    # 标记当前处理的部分
    current_section = None
    
    for line in lines:
        line = line.strip()
        if "部分积极情感词:" in line:
            current_section = "positive"
            continue
        elif "部分消极情感词:" in line:
            current_section = "negative"
            continue
        elif line.startswith("-") or line.startswith("=") or not line:
            continue
            
        # 解析词汇和分值
        if current_section and ":" in line:
            try:
                word, score_str = line.split(":")
                word = word.strip()
                score = int(score_str.strip())
                
                if current_section == "positive":
                    positive_words.append((word, score))
                elif current_section == "negative":
                    negative_words.append((word, score))
            except:
                continue
    
    return positive_words, negative_words

def main():
    print("正在生成《红楼梦》情感词云图...")
    
    # 从情感分析结果文件中加载情感词汇
    try:
        positive_words, negative_words = load_sentiment_words_from_result("sentiment_analysis_result.txt")
        print(f"加载了 {len(positive_words)} 个积极情感词")
        print(f"加载了 {len(negative_words)} 个消极情感词")
    except Exception as e:
        print(f"加载情感词汇时出错: {e}")
        # 使用默认的情感词汇
        positive_words = [("美好", 1), ("快乐", 1), ("喜悦", 1), ("高兴", 1), ("欢乐", 1), 
                         ("愉快", 1), ("欣喜", 1), ("愉悦", 1), ("开心", 1), ("得意", 1),
                         ("满意", 1), ("称赞", 1), ("赞美", 1), ("优秀", 1), ("出色", 1),
                         ("杰出", 1), ("卓越", 1), ("精彩", 1), ("完美", 1), ("美妙", 1)]
        
        negative_words = [("悲伤", -1), ("痛苦", -1), ("伤心", -1), ("难过", -1), ("忧伤", -1),
                         ("沮丧", -1), ("失望", -1), ("绝望", -1), ("悲哀", -1), ("愤怒", -1),
                         ("气愤", -1), ("生气", -1), ("怨恨", -1), ("仇恨", -1), ("厌恶", -1),
                         ("讨厌", -1), ("烦躁", -1), ("焦虑", -1), ("担心", -1), ("恐惧", -1)]
        
        print("使用默认情感词汇列表")
    
    # 生成积极情感词云图
    if positive_words:
        generate_sentiment_wordcloud(positive_words, "positive", "《红楼梦》积极情感词云图")
    else:
        print("没有积极情感词汇")
    
    # 生成消极情感词云图
    if negative_words:
        generate_sentiment_wordcloud(negative_words, "negative", "《红楼梦》消极情感词云图")
    else:
        print("没有消极情感词汇")
    
    print("词云图生成完成！")

if __name__ == "__main__":
    main()